import os
import json
import random
import cv2
import yaml
import numpy as np
from PIL import Image, ImageDraw, ImageFont

class YOLO_Dataset_Creator:
    def __init__(self, img_path, label_path, save_path, train_ratio=0.85, cls_id=4, seed=42):
        self.img_path, self.labels, self.data = img_path, label_path, save_path
        self.train = train_ratio
        self.cls = None
        self.cls_id = cls_id
        self.seed = seed

    def dataset_main(self):
        # 读取图片信息划分train和val集
        self.tr_name, self.val_name = self.divide_dataset(self.labels)
        # 根据json中类别信息生成配置yaml文件
        self.cls = self.generate_yaml(self.labels)
        # 生成yolo的txt训练数据集
        self.dataset_create(self.data)

    def divide_dataset(self, json_path):
        # 根据图片获取所有文件名信息
        total_file_list = []
        for file in os.listdir(json_path):
            if file.lower().split('.')[-1] in ['json']:
                base = file.split('.')[0]
                total_file_list.append(base)
        # 随机打乱后按比例生成训练train和验证val
        random.seed(self.seed)
        random.shuffle(total_file_list)
        length = len(total_file_list)
        tr_file_list = [tr for tr in total_file_list[:int(self.train * length)]]
        val_file_list = [te for te in total_file_list[int(self.train * length):]]
        return tr_file_list, val_file_list

    def statis_info(self, json_path):
        cls_dict = dict()
        for file in os.listdir(json_path):
            if file.lower().split('.')[-1] in ['json']:
                jsondir = os.path.join(json_path, file)
                with open(jsondir, 'r', encoding='utf-8') as f:
                    box_cls_list = json.load(f)
                for box_cls in box_cls_list:
                    cls = box_cls[self.cls_id]
                    if str(cls) not in cls_dict.keys():
                        cls_dict[str(cls)] = 1
                    else:
                        cls_dict[str(cls)] +=1
        cls_dictdir = os.path.join(self.data, 'cls_freq.json')
        with open(cls_dictdir, 'w') as f:
            json.dump(cls_dict, f)
        return cls_dict

    def generate_yaml(self, json_path):
        # 获得类别频率字典
        cls_dict = self.statis_info(json_path)
        # 生成yaml配置文件
        sorted_cls = sorted(cls_dict, key=cls_dict.get, reverse=True)
        names_dict = {}
        clses_dict = {}
        for id, c in enumerate(sorted_cls):
            names_dict[id] = c
            clses_dict[c] = id
        yaml_dict = {"path":self.data,
                     "train":"images/train",
                     "val":"images/val",
                     "names":names_dict}
        yaml_savedir = os.path.join(self.data, 'HP_Data.yaml')
        with open(yaml_savedir, "w") as f:
            yaml.dump(yaml_dict,f)
        print('yaml success')
        return clses_dict

    def dataset_create(self,data_path):
        # 必要子文件生成
        tr_img = os.path.join(data_path, 'images/train')
        va_img = os.path.join(data_path, 'images/val')
        tr_lab = os.path.join(data_path, 'labels/train')
        va_lab = os.path.join(data_path, 'labels/val')
        # 创建文件夹
        for file in [tr_img, va_img, tr_lab, va_lab]:
            os.makedirs(file, exist_ok=True)
        # train和val
        self._train(tr_img, tr_lab)
        self._val(va_img, va_lab)

    def _train(self, tr_img, tr_lab):
        # 生产训练集train
        for name in self.tr_name:
            print(name, 'start')
            # 图片复制保存移动
            imgsave = os.path.join(tr_img, name+'.jpg')
            jpgdir = os.path.join(self.img_path, name+'.jpg')
            # 标签
            txtsave = os.path.join(tr_lab, name+'.txt')
            jsondir = os.path.join(self.labels, name+'.json')
            self.label_create(jsondir, txtsave, name, jpgdir, imgsave)

    def _val(self, va_img, va_lab):
        # 生产验证集val
        for name in self.val_name:
            # 图片复制保存移动
            imgsave = os.path.join(va_img, name + '.jpg')
            jpgdir = os.path.join(self.img_path, name + '.jpg')
            # 标签
            txtsave = os.path.join(va_lab, name + '.txt')
            jsondir = os.path.join(self.labels, name + '.json')
            self.label_create(jsondir, txtsave, name, jpgdir, imgsave)

    def label_create(self, jsondir, txtsave, name, jpgdir, imgsave):
        # 图片信息
        img = cv2.imread(jpgdir)
        if img is None:
            return print(name, 'jpg is empty')
        height, width, _ = img.shape
        # 框信息
        with open(jsondir, 'r', encoding='utf-8') as f:
            box_list = json.load(f)

        box_str_list = []
        for box_cls in box_list:
            box = box_cls[:4]
            conv_box = self.normalize((width, height), box)
            text = box_cls[self.cls_id]
            cls = self.cls[str(text)]
            conv_box.insert(0, int(cls))
            box_str = " ".join(str(item) for item in conv_box)+'\n'
            box_str_list.append(box_str)
        if box_str_list!=[]:
            with open(txtsave, "w", encoding="utf-8") as f:
                f.writelines(box_str_list)

        if os.path.exists(txtsave):
            cv2.imwrite(imgsave, img)

    def normalize(self, size, box):  # size:(原图w,原图h) , box:(xmin,xmax,ymin,ymax)
        # 锚框归一化
        dw = 1. / size[0]  # 1/w
        dh = 1. / size[1]  # 1/h
        x = (box[0] + box[2]) / 2.0  # 物体在图中的中心点x坐标
        y = (box[1] + box[3]) / 2.0  # 物体在图中的中心点y坐标
        w = box[2] - box[0]  # 物体实际像素宽度
        h = box[3] - box[1]  # 物体实际像素高度
        x = x * dw  # 物体中心点x的坐标比(相当于 x/原图w)
        w = w * dw  # 物体宽度的宽度比(相当于 w/原图w)
        y = y * dh  # 物体中心点y的坐标比(相当于 y/原图h)
        h = h * dh  # 物体宽度的宽度比(相当于 h/原图h)
        return [x, y, w, h]  # 返回 相对于原图的物体中心点的x坐标比,y坐标比,宽度比,高度比,取值范围[0-1]

    def denormalize(self, size, normalized_box):  # size: (原图w, 原图h), normalized_box: [x, y, w, h]
        # 提取原图的宽度和高度
        w, h = size
        # 将中心点坐标和宽高还原为原图的像素坐标
        x_center = normalized_box[0] * w
        y_center = normalized_box[1] * h
        box_width = normalized_box[2] * w
        box_height = normalized_box[3] * h
        # 计算还原后的边界框坐标
        xmin = x_center - box_width / 2.0
        xmax = x_center + box_width / 2.0
        ymin = y_center - box_height / 2.0
        ymax = y_center + box_height / 2.0
        return [int(xmin), int(ymin), int(xmax), int(ymax)]  # 返回还原后的边界框坐标

    def visual_json_main(self, visfile):
        for file in os.listdir(self.img_path):
            base = file.split('.')[0]
            jpgdir = os.path.join(self.img_path, file)
            img = cv2.imread(jpgdir)
            if img is None:
                return print('img is empty')
            jsondir = os.path.join(self.labels, base+'.json')
            with open(jsondir, 'r', encoding='utf-8') as f:
                boxes_list = json.load(f)
            box_list = []
            for box_ in boxes_list:
                box = box_[:4]
                text = box_[self.cls_id]
                box.append(text)
                box_list.append(box)
            visdir = os.path.join(visfile, file)
            vis = self.visual_word(box_list, img ,(0,255,0))
            cv2.imwrite(visdir, vis)
            print(file,' success')

    def visual_box(self, box_list, img, color):
        for idx, box in enumerate(box_list):
            if len(box) == 4:
                l, t, r, b = box
                # 图片画框
                cv2.rectangle(img, (int(l), int(t)), (int(r), int(b)), color, thickness=2, lineType=cv2.LINE_AA)
            elif len(box) == 8:
                pts = np.array(box, np.int32)
                pts = pts.reshape((-1, 1, 2))
                cv2.polylines(img, [pts], isClosed=True, color=color, thickness=2)
        return img

    def visual_word(self, word_list, img, color):
        font_path = "C:\Windows\Fonts\SimHei.ttf"
        font_size = 45
        for idx, box in enumerate(word_list):
            l, t, r, b, word = box
            # 图片画框
            cv2.rectangle(img, (l, t), (r, b), color, thickness=2, lineType=cv2.LINE_AA)
            caption = f"{word}"
            pil_img = Image.new('RGB', (0, 0))
            draw = ImageDraw.Draw(pil_img)
            text_size = draw.textbbox((0, 0), caption, font=ImageFont.truetype(font_path, font_size))
            text_width = text_size[2] - text_size[0]
            text_height = text_size[3] - text_size[1]
            # 使用 OpenCV 画文本背景框
            cv2.rectangle(img, (r, t), (r + text_width, t + text_height), color, -1)
            # 使用 PIL 绘制中文标签
            img = self.put_chinese_text(img, caption, (r, t), font_size, font_path, (0, 0, 0))
        return img

    def put_chinese_text(self, img, text, position, font_size, font_path, text_color):
        # 创建一个 PIL 图像
        pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(pil_img)
        # 加载字体
        font = ImageFont.truetype(font_path, font_size)
        # 绘制文本
        draw.text(position, text, font=font, fill=text_color)
        # 将 PIL 图像转换回 OpenCV 格式
        img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
        return img

# 逆向工程
if __name__=='__main__':
    jpgfile = './example/origin/img/'
    jsonfile = './example/origin/json/'
    visfile = './example/origin/vis/'
    # 读取配置类别
    with open('./example/official/coco8.yaml', 'r', encoding='utf-8')as f:
        cfg = yaml.safe_load(f)
    cls_list = cfg['names']

    yolo = YOLO_Dataset_Creator(None, None, None)
    for file in os.listdir('./example/official'):
        if file.lower().endswith('.txt'):
            txtdir = os.path.join('./example/official', file)
            with open(txtdir, 'r') as f:
                txt = f.readlines()
            base = file.split('.')[0]
            jpgdir = os.path.join(jpgfile, base+'.jpg')
            img = cv2.imread(jpgdir)
            height,width,_ = img.shape

            box_list = []
            for t in txt:
                cls_box = t[:-1]
                bbox = cls_box.split(' ')
                cls = int(bbox[0])
                cls_text = cls_list[cls]
                box = [float(bbox[1]), float(bbox[2]), float(bbox[3]), float(bbox[4])]
                normal_box = yolo.denormalize((width, height), box)
                normal_box.append(cls_text)
                box_list.append(normal_box)
            jsondir = os.path.join(jsonfile, base+'.json')
            with open(jsondir, 'w') as f:
                json.dump(box_list, f)

            vis = yolo.visual_word(box_list, img, (0,255,0))
            visdir = os.path.join(visfile, base+'.jpg')
            cv2.imwrite(visdir, vis)
